import os
import pandas as pd
import numpy as np
import csv

#
population_csv = pd.read_csv("data/population/population.csv")
standard = pd.read_csv("csv_results/standard.csv")
standard = dict(zip(standard.values[:, 0].tolist(), standard.values[:, 1].tolist()))
output = {}
max_y = -0x3f3f3f3f
min_y = 0x3f3f3f3f
fid = []
num = 0
for filename in os.listdir("data/deleted-and-transposed"):
    csv_data = pd.read_csv("data/deleted-and-transposed/" + filename).values
    year = csv_data[:, 0]
    max_y = int(max(np.max(year), max_y))
    min_y = int(min(np.min(year), min_y))
    num += 1
print(max_y)
print(min_y)
# 时间行，种类数列。
table = np.zeros((max_y - min_y + 1, num), dtype=np.float32)
table.fill(np.nan)
for i, filename in enumerate(os.listdir("data/deleted-and-transposed")):
    csv_file = pd.read_csv("data/deleted-and-transposed/" + filename)
    csv_data = csv_file.values
    file_id = filename[:filename.rfind(".")]
    countries = csv_file.columns.tolist()
    if "ISO" in countries:
        countries.remove("ISO")
    population_num_all = population_csv[countries].values  # 人口数量
    population_year = population_csv["year"].values
    data = csv_data[:, 1:]
    year = np.int32(csv_data[:, 0])
    population_num = []
    for y in year:
        population_this_year = population_num_all[population_year == y]
        assert population_this_year.shape[0] == 1
        population_num.append(population_this_year[0])
    data_mean = np.zeros_like(year, dtype=float)

    for lineid in range(data.shape[0]):
        if year[lineid] == 2022:
            print("haha")
        data_line = data[lineid]
        population_line = population_num[lineid]
        mask = np.logical_and(np.logical_not(np.isnan(data_line)), np.logical_not(np.isnan(population_line)))
        if np.sum(mask) == 0:
            data_mean[lineid] = np.nan
        else:
            data_mean[lineid] = np.sum(data_line[mask] * population_line[mask]) / np.sum(population_line[mask])
    print(data_mean[-1], year[-1])
    print("years = "+str(year - min_y))
    table[year - min_y, i] = data_mean
    print("table"+str(table[-1]))
    fid.append(file_id)
# table1 = np.zeros((max_y - min_y + 1, 18), dtype=np.float32)
# cnt1 = np.zeros((max_y - min_y + 1, 18), dtype=np.int32)
# cnt1[:, 0] = 1
# table1[:, 0] = np.arange(min_y, max_y + 1)
# for iid, ffid in enumerate(fid):
#     table1[:, ffid] += table[:, iid]
#     cnt1[:, ffid] += 1
# table1 /= cnt1
# table1_list = table1.tolist()
# for ls in table1_list:
#     ls[0] = int(ls[0])
#
# for ls in table1_list:
#     for iid, item in enumerate(ls):
#         if np.isnan(item):
#             ls[iid] = ''
with open("csv_results/means.csv", "w", newline='', encoding="utf-8") as f:
    writer = csv.writer(f)
    writer.writerow(["year"] + fid)
    for i, line in enumerate(table):
        writer.writerow([min_y + i] + line.tolist())
